import json
import time
import traceback
from typing import Any, Callable
from .tool import list_tool_descriptors, get_tool_function
from .api import chat_completion_create, chat_completion_create_streamed

def ask_question(model: str, messages: list[dict], options: dict[str, Any], callbacks: dict[str, Callable]) -> str:
    answer = None
    while answer is None:
        tools = list_tool_descriptors(options.get('tool_categories', ()))
        t0 = time.time()
        kwargs = dict(
            model=model,
            messages=messages,  # type: ignore
            tools=tools,
            max_tokens=options.get('max_tokens'),
            temperature=options.get('temperature'),
        )
        append = True
        if options.get('streamed', False):
            def stream_callback(completion):
                nonlocal append
                message = completion.choices[0].message
                if append:
                    messages.append(message)
                else:
                    messages[-1] = message
                if 'message' in callbacks:
                    dt = time.time() - t0
                    message_dumped = message.model_dump()
                    if message_dumped.get('tool_calls', None) is None:
                        message_dumped.pop('tool_calls', None)
                    message_dumped.pop('function_call', None)
                    callbacks['message'](message_dumped, dt, append, False)
                append = False
                if 'answer' in callbacks:
                    if message.content:
                        callbacks['answer'](message.content, False)
            completion = chat_completion_create_streamed(
                callback=stream_callback,
                **kwargs,
            )
        else:
            completion = chat_completion_create(
                **kwargs,
            )
        dt = time.time() - t0
        choice = completion.choices[0]
        message = choice.message
        if choice.finish_reason == 'length':
            print('WARNING: Output truncated')
        message_dumped = message.model_dump()
        if message_dumped.get('tool_calls', None) is None:
            message_dumped.pop('tool_calls', None)
        message_dumped.pop('function_call', None)
        if append:
            messages.append(message_dumped)
        else:
            messages[-1] = message_dumped
        if 'message' in callbacks:
            callbacks['message'](message_dumped, dt, append, True)
        append = False
        if message.content is not None:
            if 'answer' in callbacks:
                callbacks['answer'](message.content, True)
            if message.tool_calls is None:
                answer = message.content
        if message.tool_calls is not None:
            for tool_call in message.tool_calls:
                function_name = tool_call.function.name
                t0 = time.time()
                function_to_call = get_tool_function(function_name)
                if function_to_call is None:
                    function_response = {'success': False, 'error': f'Function {function_name} not found'}
                else:
                    try:
                        function_args = json.loads(tool_call.function.arguments)
                        if 'check' in callbacks and not callbacks['check'](function_name, function_args):
                            function_response = {'success': False, 'message': 'User refused this action'}
                            if answer is None:
                                answer = ''
                        else:
                            function_response = function_to_call(**function_args)
                    except:
                        traceback.print_exc()
                        function_response = {'success': False, 'traceback': traceback.format_exc()}
                if not isinstance(function_response, str) and function_response is not None:
                    try:
                        function_response = json.dumps(function_response, ensure_ascii=False, separators=(',', ':'))
                    except:
                        function_response = str(function_response)
                dt = time.time() - t0
                if function_response is None:
                    function_response = '{"success": true}'
                    if answer is None:
                        answer = ''
                message = {
                    'tool_call_id': tool_call.id,
                    'role': 'tool',
                    'name': function_name,
                    'content': function_response,
                }
                if 'message' in callbacks:
                    callbacks['message'](message, dt, True, True)
                messages.append(message)
    return answer
